#include <cuda_runtime.h>
#include <cuda_fp16.h> 
#include "conv2d.h"

extern "C" __global__ void implicit_gemm_v0(mykernelParamType param)
{ 
    // 计算当前线程处理的输出位置和通道
    int OhOw = blockIdx.x * blockDim.x + threadIdx.x;  // 输出空间位置(oh*ow)
    int k = blockIdx.y * blockDim.y + threadIdx.y;     // 输出通道
    int n = blockIdx.z * blockDim.z + threadIdx.z;     // batch索引
    
    // 边界检查
    if(OhOw >= param.Oh*param.Ow || k >= param.k || n >= param.n)
        return;

    // 分解输出位置为oh和ow
    int oh = OhOw / param.Ow;  // 输出高度坐标
    int ow = OhOw % param.Ow;  // 输出宽度坐标
    
    int input_addr, weight_addr, output_addr;
    float sum = 0.0f;
    
    // 计算输出地址 [n,k,oh,ow] 布局
    output_addr = n * param.k * param.Oh * param.Ow + 
                 k * param.Oh * param.Ow + 
                 oh * param.Ow + ow;

    // 遍历卷积核的空间维度
    for(int r = 0; r < param.r; r++) {
        for(int s = 0; s < param.s; s++) {
            // 计算输入位置（考虑stride和padding）
            int ih = oh * param.u + r - param.p;
            int iw = ow * param.v + s - param.q;

            // 检查输入位置是否有效
            if (ih >= 0 && ih < param.h && iw >= 0 && iw < param.w) {
                // 遍历输入通道
                for(int c = 0; c < param.c; c++) {
                    // 计算输入数据地址 [n,c,h,w] 布局
                    input_addr = n * param.c * param.h * param.w +c * param.h * param.w +ih * param.w + iw;
                    weight_addr = k * param.c * param.r * param.s +c * param.r * param.s +r * param.s + s;

                    sum += param.pin[input_addr] * param.pweight[weight_addr];
                }
            }
        }
    }
    // 写入结果到全局内存
    param.pout[output_addr] = sum;
}

void launch_implicit_gemm_v0(unsigned int outh, unsigned int outw, unsigned int k, unsigned int n, mykernelParamType* param) {
    // 计算网格维度
    int blockx = (param->Oh * param->Ow + 15) / 16;  // 输出空间维度分块数
    int blocky = (param->k + 15) / 16;              // 输出通道维度分块数
    int blockz = (param->n + 0) / 1;                // batch维度分块数
    
    // 设置线程块维度
    int threadx = 16;  // x方向16线程
    int thready = 16;  // y方向16线程
    int threadz = 1;   // z方向1线程
    
    dim3 block(threadx, thready, threadz);  // 线程块维度
    dim3 grid(blockx, blocky, blockz);      // 网格维度
    
    // 启动核函数
    implicit_gemm_v0<<<grid, block>>>(*param);
}